#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import sys
import subprocess
import os
import re
import shutil
import tempfile
import time
import argparse
from typing import List, Dict, Optional, Tuple

# 版本信息
KM_VERSION = "1.2.2-8"

class CLILogger:
    """CLI日志类"""
    @staticmethod
    def add_log(message: str):
        """添加日志信息"""
        print(f"[LOG] {message}")
    
    @staticmethod
    def add_error(message: str):
        """添加错误日志"""
        print(f"[ERROR] {message}")
    
    @staticmethod
    def add_warning(message: str):
        """添加警告日志"""
        print(f"[WARNING] {message}")

class PackageManager:
    """包管理器类"""
    def __init__(self):
        self.package_manager = self._detect_package_manager()
        
    def _detect_package_manager(self):
        """通过检测可用的包管理器来确定系统类型"""
        # 检测 apt (Debian/Ubuntu 系列)
        if shutil.which('apt-get'):
            return 'apt'
        else:
            raise RuntimeError("不支持的包管理器，仅支持Ubuntu/Debian系统")
            
    def get_package_names(self, kernel_version, action='install'):
        """根据包管理器和操作类型生成相关的包名列表"""
        packages = []
        
        if self.package_manager == 'apt':
            if action == 'install':
                packages.extend([
                    f'linux-image-{kernel_version}',
                    f'linux-headers-{kernel_version}',
                ])
            else:  # remove
                packages.extend([
                    f'linux-image-{kernel_version}',
                    f'linux-headers-{kernel_version}',
                ])
        else:
            raise RuntimeError(f"不支持的包管理器: {self.package_manager}")
            
        return packages
    
    def install_packages(self, package_names, log_window=None):
        """安装多个包"""
        self._install_package_list(package_names, log_window)
        return package_names
    
    def remove_packages(self, package_names, log_window=None):
        """移除多个包"""
        self._remove_package_list(package_names, log_window)
    
    def _install_package_list(self, package_names, log_window=None):
        """安装包列表"""
        try:
            if self.package_manager == 'apt':
                if os.geteuid() != 0:
                    try:
                        if log_window:
                            log_window.add_log("使用 pkexec 执行安装命令...")
                        result = subprocess.Popen(['pkexec', 'apt-get', 'install', '-y'] + package_names, 
                                                stdout=subprocess.PIPE, stderr=subprocess.PIPE, 
                                                universal_newlines=True)
                        # 实时读取输出
                        for line in iter(result.stdout.readline, ''):
                            if log_window:
                                log_window.add_log(line.rstrip())
                        for line in iter(result.stderr.readline, ''):
                            if log_window:
                                log_window.add_log("错误: " + line.rstrip())
                        result.wait()
                        if result.returncode != 0:
                            raise subprocess.CalledProcessError(result.returncode, result.args)
                    except (subprocess.CalledProcessError, FileNotFoundError):
                        try:
                            if log_window:
                                log_window.add_log("使用 sudo 执行安装命令...")
                            result = subprocess.Popen(['sudo', 'apt-get', 'install', '-y'] + package_names,
                                                    stdout=subprocess.PIPE, stderr=subprocess.PIPE, 
                                                    universal_newlines=True)
                            # 实时读取输出
                            for line in iter(result.stdout.readline, ''):
                                if log_window:
                                    log_window.add_log(line.rstrip())
                            for line in iter(result.stderr.readline, ''):
                                if log_window:
                                    log_window.add_log("错误: " + line.rstrip())
                            result.wait()
                            if result.returncode != 0:
                                raise subprocess.CalledProcessError(result.returncode, result.args)
                        except (subprocess.CalledProcessError, FileNotFoundError):
                            raise PermissionError("无法安装包")
                else:
                    if log_window:
                        log_window.add_log("以 root 权限执行安装命令...")
                    result = subprocess.Popen(['apt-get', 'install', '-y'] + package_names, 
                                            stdout=subprocess.PIPE, stderr=subprocess.PIPE, 
                                            universal_newlines=True)
                    # 实时读取输出
                    for line in iter(result.stdout.readline, ''):
                        if log_window:
                            log_window.add_log(line.rstrip())
                    for line in iter(result.stderr.readline, ''):
                        if log_window:
                            log_window.add_log("错误: " + line.rstrip())
                    result.wait()
                    if result.returncode != 0:
                        raise subprocess.CalledProcessError(result.returncode, result.args)
                        
        except Exception as e:
            raise RuntimeError(f"安装包失败: {str(e)}")
    
    def _remove_package_list(self, package_names, log_window=None):
        """移除包列表"""
        try:
            if self.package_manager == 'apt':
                if os.geteuid() != 0:
                    try:
                        if log_window:
                            log_window.add_log("使用 pkexec 执行移除命令...")
                        result = subprocess.Popen(['pkexec', 'apt-get', 'remove', '-y'] + package_names, 
                                                stdout=subprocess.PIPE, stderr=subprocess.PIPE, 
                                                universal_newlines=True)
                        # 实时读取输出
                        for line in iter(result.stdout.readline, ''):
                            if log_window:
                                log_window.add_log(line.rstrip())
                        for line in iter(result.stderr.readline, ''):
                            if log_window:
                                log_window.add_log("错误: " + line.rstrip())
                        result.wait()
                        if result.returncode != 0:
                            raise subprocess.CalledProcessError(result.returncode, result.args)
                    except (subprocess.CalledProcessError, FileNotFoundError):
                        try:
                            if log_window:
                                log_window.add_log("使用 sudo 执行移除命令...")
                            result = subprocess.Popen(['sudo', 'apt-get', 'remove', '-y'] + package_names,
                                                    stdout=subprocess.PIPE, stderr=subprocess.PIPE, 
                                                    universal_newlines=True)
                            # 实时读取输出
                            for line in iter(result.stdout.readline, ''):
                                if log_window:
                                    log_window.add_log(line.rstrip())
                            for line in iter(result.stderr.readline, ''):
                                if log_window:
                                    log_window.add_log("错误: " + line.rstrip())
                            result.wait()
                            if result.returncode != 0:
                                raise subprocess.CalledProcessError(result.returncode, result.args)
                        except (subprocess.CalledProcessError, FileNotFoundError):
                            raise PermissionError("无法移除包")
                else:
                    if log_window:
                        log_window.add_log("以 root 权限执行移除命令...")
                    result = subprocess.Popen(['apt-get', 'remove', '-y'] + package_names, 
                                            stdout=subprocess.PIPE, stderr=subprocess.PIPE, 
                                            universal_newlines=True)
                    # 实时读取输出
                    for line in iter(result.stdout.readline, ''):
                        if log_window:
                            log_window.add_log(line.rstrip())
                    for line in iter(result.stderr.readline, ''):
                        if log_window:
                            log_window.add_log("错误: " + line.rstrip())
                    result.wait()
                    if result.returncode != 0:
                        raise subprocess.CalledProcessError(result.returncode, result.args)
                        
        except Exception as e:
            raise RuntimeError(f"移除包失败: {str(e)}")


class InstallRecordManager:
    RECORD_FILE = os.path.expanduser("~/.config/kernel-management/install-records.log")
    
    @classmethod
    def add_record(cls, kernel_version, packages):
        """添加安装记录"""
        try:
            os.makedirs(os.path.dirname(cls.RECORD_FILE), exist_ok=True)
            with open(cls.RECORD_FILE, 'a') as f:
                f.write(f"{kernel_version}:{','.join(packages)}\n")
        except PermissionError:
            print(f"权限不足，无法写入安装记录文件: {cls.RECORD_FILE}")
        except Exception as e:
            print(f"写入安装记录失败: {e}")
    
    @classmethod
    def remove_record(cls, kernel_version):
        """移除安装记录"""
        try:
            if not os.path.exists(cls.RECORD_FILE):
                return
            
            lines = []
            with open(cls.RECORD_FILE, 'r') as f:
                lines = f.readlines()
            
            with open(cls.RECORD_FILE, 'w') as f:
                for line in lines:
                    if not line.startswith(f"{kernel_version}:"):
                        f.write(line)
        except Exception:
            pass  # 忽略记录失败
    

    @classmethod
    def load_records(cls):
        """加载安装记录"""
        records = {}
        try:
            if os.path.exists(cls.RECORD_FILE):
                with open(cls.RECORD_FILE, 'r') as f:
                    for line in f:
                        if ':' in line:
                            kernel_version, packages = line.strip().split(':', 1)
                            records[kernel_version] = packages.split(',')
        except Exception:
            pass  # 忽略记录失败
        return records

class DefaultBootManager:
    """默认启动管理器"""
    DEFAULT_BOOT_FILE = os.path.expanduser("~/.config/kernel-management/default-kernel.log")
    NEXT_BOOT_FILE = "/tmp/kernel-management/next-boot-kernel.log"

    
    @classmethod
    def save_default_boot(cls, kernel_version):
        """保存默认启动内核记录"""
        try:
            os.makedirs(os.path.dirname(cls.DEFAULT_BOOT_FILE), exist_ok=True)
            with open(cls.DEFAULT_BOOT_FILE, 'w') as f:
                f.write(kernel_version)
        except Exception:
            pass  # 忽略记录失败
    
    @classmethod
    def save_next_boot(cls, kernel_version):
        """保存下次启动内核记录"""
        try:
            os.makedirs(os.path.dirname(cls.NEXT_BOOT_FILE), exist_ok=True)
            with open(cls.NEXT_BOOT_FILE, 'w') as f:
                f.write(kernel_version)
        except Exception:
            pass  # 忽略记录失败
    
    @classmethod
    def load_default_boot(cls):
        """加载默认启动内核记录"""
        try:
            if os.path.exists(cls.DEFAULT_BOOT_FILE):
                with open(cls.DEFAULT_BOOT_FILE, 'r') as f:
                    return f.read().strip()
        except Exception:
            pass
        return None
    
    @classmethod
    def load_next_boot(cls):
        """加载下次启动内核记录"""
        try:
            if os.path.exists(cls.NEXT_BOOT_FILE):
                with open(cls.NEXT_BOOT_FILE, 'r') as f:
                    return f.read().strip()
        except Exception:
            pass
        return None

class BootManager:
    GRUB_CONFIG_PATHS = [
        "/boot/grub/grub.cfg",
        "/boot/grub2/grub.cfg"
    ]
    
    GRUB_DEFAULT_PATH = "/etc/default/grub"
    GRUB_ENV_PATH = "/boot/grub/grubenv"
    
    def __init__(self, cli_instance):
        self.cli_instance = cli_instance  # 保存CLI实例的引用
        self.grub_cfg_path = self._detect_grub_config()
        
    def _detect_grub_config(self):
        """检测GRUB配置文件路径"""
        for path in self.GRUB_CONFIG_PATHS:
            if os.path.exists(path):
                return path
        return self.GRUB_CONFIG_PATHS[0]  # 默认返回第一个
    
    def get_boot_entries(self):
        """获取启动条目列表"""
        return self.get_grub_entries()
    
    def get_grub_entries(self):
        """获取GRUB启动条目"""
        entries = []
        
        try:
            # 读取GRUB配置文件
            if os.geteuid() != 0:
                try:
                    result = subprocess.run(['pkexec', 'cat', self.grub_cfg_path], 
                                        capture_output=True, text=True, check=True)
                    content = result.stdout
                except (subprocess.CalledProcessError, FileNotFoundError):
                    try:
                        result = subprocess.run(['sudo', 'cat', self.grub_cfg_path], 
                                            capture_output=True, text=True, check=True)
                        content = result.stdout
                    except (subprocess.CalledProcessError, FileNotFoundError):
                        raise PermissionError("无法读取GRUB配置文件")
            else:
                with open(self.grub_cfg_path, 'r') as f:
                    content = f.read()
            
            # 解析结构
            current_submenus = []  # 跟踪当前的submenu路径
            lines = content.split('\n')
            
            i = 0
            while i < len(lines):
                line = lines[i]
                stripped_line = line.strip()
                
                # 跳过空行
                if not stripped_line:
                    i += 1
                    continue
                
                # 检查submenu
                submenu_match = re.match(r"submenu\s+'([^']+)'", stripped_line)
                if submenu_match:
                    submenu_title = submenu_match.group(1)
                    # 检查是否在行首
                    is_at_line_start = line == stripped_line
                    if is_at_line_start:
                        # 顶级submenu，重置路径
                        current_submenus = [submenu_title]
                    else:
                        # 嵌套submenu，添加到当前路径
                        # 简化处理：我们假设正常的GRUB结构
                        if current_submenus:
                            # 保持当前路径不变，因为我们主要关心menuentry的位置
                            pass
                    i += 1
                    continue
                
                # 检查menuentry
                menuentry_match = re.match(r"menuentry\s+'([^']+)'", stripped_line)
                if menuentry_match:
                    title = menuentry_match.group(1)
                    # 查找内核版本
                    version = None
                    j = i + 1
                    while j < len(lines) and lines[j].strip() != '}':
                        linux_match = re.search(r'linux.*?/boot/vmlinuz-(\S+)', lines[j])
                        if linux_match:
                            version = linux_match.group(1)
                            break
                        j += 1
                    
                    # 如果通过linux命令未找到版本，尝试从标题中提取
                    if not version:
                        version = self._extract_kernel_version_from_title(title)
                    
                    # 判断是否为子菜单项：检查menuentry是否在行首
                    is_at_line_start = line == stripped_line
                    is_submenu_entry = not is_at_line_start  # 不在行首的就是子菜单项
                    
                    # 构建完整路径
                    if is_submenu_entry and current_submenus:
                        # 如果是子菜单项且有submenu上下文，则构建完整路径
                        full_path = ">".join(current_submenus + [title])
                    else:
                        # 否则只使用标题
                        full_path = title
                        # 如果是顶级menuentry，重置submenu上下文
                        if is_at_line_start:
                            current_submenus = []
                    
                    entries.append({
                        'title': title,
                        'version': version,
                        'identifier': full_path,
                        'is_submenu_entry': is_submenu_entry
                    })
                
                i += 1
                        
        except Exception as e:
            print(f"解析GRUB条目失败: {e}")
            raise RuntimeError(f"无法解析GRUB启动条目: {e}")
        
        return entries

    def _extract_kernel_version_from_title(self, title):
        """
        从标题中提取内核版本信息
        特别针对深度操作系统等自定义格式的GRUB菜单
        """
        # 常见的内核版本模式
        patterns = [
            r'linux\s+([\d\.\-\w]+)',  # 匹配 "linux X.X.X-..." 格式
            r'linux\s+version\s+([\d\.\-\w]+)',  # 匹配 "linux version X.X.X-..." 格式
            r'kernel\s+([\d\.\-\w]+)',  # 匹配 "kernel X.X.X-..." 格式
            r'version\s+([\d\.\-\w]+)',  # 匹配 "version X.X.X-..." 格式
            r'(\d+\.\d+\.\d+[-\w]*)',  # 匹配纯版本号格式如 "X.X.X-..."
            r'Linux\s+([\d\.\-\w]+)',  # 匹配 "Linux X.X.X-..." 格式
            r'with\s+Linux\s+([\d\.\-\w]+)',  # 匹配 "with Linux X.X.X-..." 格式
        ]
        
        # 针对深度操作系统的特殊处理
        # 匹配 "Linux X.X.X-amd64-desktop-rolling" 格式
        deepin_pattern = r'Linux\s+(\d+\.\d+\.\d+[-\w]*)'
        match = re.search(deepin_pattern, title)
        if match:
            return match.group(1)
        
        # 尝试其他通用模式
        for pattern in patterns:
            match = re.search(pattern, title, re.IGNORECASE)
            if match:
                return match.group(1)
        
        # 如果以上都没匹配到，返回None
        return None

    
    def set_default_entry(self, identifier):
        """设置默认启动条目"""
        self.set_default_kernel_with_recommended_method(identifier, permanent=True)
    
    def set_next_boot_entry(self, identifier):
        """设置下次启动条目"""
        self.set_default_kernel_with_recommended_method(identifier, permanent=False)


    def set_default_kernel_with_recommended_method(self, identifier=None, permanent=True):
        """
        使用GRUB推荐方式设置默认启动内核
        现在默认使用GRUB 2.0+的方式
        """
        log_window = CLILogger()  # 使用CLI日志
        action_type = "永久" if permanent else "临时"  # 初始化 action_type
        
        try:
            # 获取选中的内核条目
            kernel_list = self.cli_instance.get_all_available_kernels()
            installed_kernels = self.cli_instance.get_installed_kernels()
            current_kernel = self.cli_instance.get_current_kernel()
            
            # 显示可用内核列表
            print("可用内核列表:")
            for i, kernel in enumerate(kernel_list):
                installed_mark = "✓" if kernel in installed_kernels else "○"
                current_mark = "[当前]" if current_kernel and kernel.replace('linux-image-', '') in current_kernel else ""
                print(f"  {i+1}. {installed_mark} {kernel} {current_mark}")
            
            # 获取用户选择
            try:
                choice = int(input("请选择要设置为默认启动的内核编号: ")) - 1
                if choice < 0 or choice >= len(kernel_list):
                    print("无效的选择")
                    return
            except ValueError:
                print("无效的输入")
                return
                
            kernel_name = kernel_list[choice]
            
            # 根据不同包管理器提取版本号
            if kernel_name.startswith('linux-image-'):
                kernel_version = kernel_name.replace('linux-image-', '')
            elif kernel_name.startswith('kernel-'):
                kernel_version = kernel_name.replace('kernel-', '')
            else:
                kernel_version = kernel_name.replace('linux-', '')
            
            # 确认设置
            confirm = input(f"确定要将 {kernel_name} 设置为{action_type}默认启动内核吗？(y/N): ")
            if confirm.lower() != 'y':
                return
                
            log_window.add_log(f"开始设置{action_type}默认启动内核: {kernel_name}")
            log_window.add_log(f"内核版本: {kernel_version}")
            
            # 获取GRUB版本
            grub_version = self.get_grub_version()
            log_window.add_log(f"检测到GRUB版本: {grub_version}")
            
            # 查找目标内核条目
            target_entry = self.find_grub_entry_for_kernel(kernel_version)
            if not target_entry:
                raise RuntimeError(f"在GRUB菜单中找不到内核 {kernel_version} 的条目")
            
            log_window.add_log(f"找到目标条目: {target_entry['title']} (版本: {target_entry['version']})")
            
            # 现在默认使用GRUB 2.0+的方式设置
            menu_path = self.construct_grub_200_plus_menu_path(kernel_version)
            log_window.add_log(f"使用GRUB 2.0+推荐方式设置默认启动项: {menu_path}")
            
            # 设置默认启动项
            if permanent:
                self._update_grub_default_config(f'GRUB_DEFAULT="{menu_path}"')
                log_window.add_log("已永久设置GRUB默认启动项")
                # 记录到本地文件
                DefaultBootManager.save_default_boot(kernel_version)
                log_window.add_log("已记录为默认启动内核")
                # 更新GRUB配置
                self.update_grub_config()
                log_window.add_log("已更新GRUB配置")
            else:
                self._set_next_boot_entry(menu_path)
                log_window.add_log("已设置下次启动项")
                # 记录到本地文件
                DefaultBootManager.save_next_boot(kernel_version)
                log_window.add_log("已记录为下次启动内核")
            
            log_window.add_log(f"{action_type}默认启动条目设置完成")
            log_window.add_log(f"已将 {kernel_name} 设置为{action_type}默认启动内核（使用GRUB推荐方式）")
            
            print(f"成功: 已将 {kernel_name} 设置为{action_type}默认启动内核")
            
        except Exception as e:
            log_window.add_log(f"设置{action_type}默认启动内核失败: {str(e)}")
            print(f"错误: 设置{action_type}默认启动内核失败: {str(e)}")

    def get_grub_version(self):
        """获取GRUB版本"""
        try:
            result = subprocess.run(['grub-install', '--version'], 
                                  capture_output=True, text=True, check=True)
            version_line = result.stdout.strip()
            # 提取版本号 (格式: "grub-install (GRUB) 2.06")
            import re
            match = re.search(r'GRUB\)\s+([\d.]+)', version_line)
            if match:
                return match.group(1)
        except (subprocess.CalledProcessError, FileNotFoundError):
            pass
        return None

    def is_grub_200_or_later(self, version_str):
        """判断GRUB版本是否为2.00或更高"""
        # 修改为始终返回True，即默认使用GRUB 2.0+的方式
        return True
        
    
    def find_grub_entry_for_kernel(self, target_kernel_version):
        """查找GRUB中对应内核的条目"""
        try:
            # 获取GRUB条目
            entries = self.get_grub_entries()
            if not entries:
                return None
            
            # 查找目标内核条目（优先选择非恢复模式的条目）
            target_entry = None
            for entry in entries:
                # 首先尝试通过版本号匹配
                if (entry['version'] and 
                    (target_kernel_version == entry['version'] or 
                    target_kernel_version in entry['version'] or 
                    entry['version'] in target_kernel_version) and
                    'recovery' not in entry['title'].lower() and 
                    '恢复' not in entry['title']):
                    target_entry = entry
                    break
                
                # 如果没有通过版本号匹配，则尝试通过标题匹配
                elif (target_kernel_version in entry['title'] and
                    'recovery' not in entry['title'].lower() and
                    '恢复' not in entry['title']):
                    target_entry = entry
                    break
            
            # 如果没找到非恢复模式条目，则查找任何匹配的条目（包括恢复模式）
            if not target_entry:
                for entry in entries:
                    if (entry['version'] and 
                        (target_kernel_version == entry['version'] or 
                        target_kernel_version in entry['version'] or 
                        entry['version'] in target_kernel_version)):
                        target_entry = entry
                        break
                    elif target_kernel_version in entry['title']:
                        target_entry = entry
                        break
            
            return target_entry
        except Exception as e:
            print(f"查找GRUB条目失败: {e}")
            return None

    def construct_grub_200_plus_menu_path(self, kernel_version):
        """
        为深度系统优化GRUB菜单路径构建
        """
        try:
            # 获取GRUB条目
            entries = self.get_grub_entries()
            
            # 针对深度系统的特殊处理
            target_entry = None
            
            # 优先查找匹配的高级选项菜单项
            for entry in entries:
                if (entry['version'] and 
                    (kernel_version == entry['version'] or 
                    kernel_version in entry['version']) and
                    'recovery' not in entry['title'].lower() and
                    '恢复' not in entry['title'] and
                    entry['is_submenu_entry']):
                    target_entry = entry
                    break
            
            # 如果没找到，尝试查找任何匹配项
            if not target_entry:
                for entry in entries:
                    if (entry['version'] and 
                        (kernel_version == entry['version'] or 
                        kernel_version in entry['version']) and
                        'recovery' not in entry['title'].lower() and
                        '恢复' not in entry['title']):
                        target_entry = entry
                        break
            
            if target_entry and 'identifier' in target_entry:
                return target_entry['identifier']
                
            raise RuntimeError(f"找不到内核版本 {kernel_version} 的GRUB条目")
                
        except Exception as e:
            raise RuntimeError(f"构建GRUB菜单路径失败: {e}")


    def get_root_uuid(self):
        """获取根分区UUID"""
        try:
            # 尝试多种方式获取根分区UUID
            commands = [
                ['lsblk', '-f', '-o', 'UUID,MOUNTPOINT'],
                ['sudo', 'lsblk', '-f', '-o', 'UUID,MOUNTPOINT'],
                ['pkexec', 'lsblk', '-f', '-o', 'UUID,MOUNTPOINT']
            ]
            
            for cmd in commands:
                try:
                    result = subprocess.run(cmd, capture_output=True, text=True, check=True)
                    lines = result.stdout.strip().split('\n')
                    for line in lines[1:]:  # 跳过标题行
                        parts = line.split()
                        if len(parts) >= 2 and '/' in parts[-1] and parts[-1].strip() == '/':
                            return parts[0]  # 返回根分区UUID
                except (subprocess.CalledProcessError, FileNotFoundError):
                    continue
                    
            # 如果上面的方法失败，尝试使用blkid
            try:
                result = subprocess.run(['blkid'], capture_output=True, text=True, check=True)
                # 查找包含"/"的行（根分区）
                import re
                for line in result.stdout.split('\n'):
                    if 'UUID=' in line and '/dev/' in line:
                        uuid_match = re.search(r'UUID="([^"]+)"', line)
                        if uuid_match and '/dev/' in line:
                            # 进一步确认是否为根分区
                            if 'TYPE="ext"' in line or 'TYPE="btrfs"' in line or 'TYPE="xfs"' in line:
                                return uuid_match.group(1)
            except (subprocess.CalledProcessError, FileNotFoundError):
                pass
                
        except Exception as e:
            print(f"获取根分区UUID失败: {e}")
        
        return None

    def _update_grub_default_config(self, new_default_line):
        """更新GRUB默认配置文件 - 安全版本"""
        try:
            # 验证输入
            if not self._is_valid_grub_default(new_default_line):
                raise ValueError("无效的GRUB默认配置")
            
            # 使用更安全的临时文件处理
            with tempfile.NamedTemporaryFile(mode='w', delete=False, prefix='grub_default_') as tmp_file:
                # 读取原配置文件
                original_lines = []
                if os.path.exists(self.GRUB_DEFAULT_PATH):
                    with open(self.GRUB_DEFAULT_PATH, 'r') as f:
                        original_lines = f.readlines()
                
                # 写入修改后的配置
                grub_default_found = False
                for line in original_lines:
                    if line.startswith('GRUB_DEFAULT='):
                        tmp_file.write(new_default_line + '\n')
                        grub_default_found = True
                    else:
                        tmp_file.write(line)
                
                # 如果原来没有GRUB_DEFAULT，则添加
                if not grub_default_found:
                    tmp_file.write(new_default_line + '\n')
                
                tmp_path = tmp_file.name
            
            try:
                # 安全的文件替换
                if os.geteuid() != 0:
                    # 尝试使用 pkexec 进行替换
                    subprocess.run(['pkexec', 'mv', tmp_path, self.GRUB_DEFAULT_PATH], check=True)
                else:
                    shutil.move(tmp_path, self.GRUB_DEFAULT_PATH)
            except Exception as e:
                os.unlink(tmp_path)  # 确保清理临时文件
                raise RuntimeError(f"更新GRUB配置失败: {str(e)}")
        
        except Exception as e:
            raise RuntimeError(f"更新GRUB配置失败: {str(e)}")

    def _is_valid_grub_default(self, line):
        """验证GRUB默认配置行是否有效"""
        # 简单验证GRUB_DEFAULT格式
        pattern = r'^GRUB_DEFAULT="[^"]*"$'
        return re.match(pattern, line) is not None

    def _set_next_boot_entry(self, menu_path):
        """设置下次启动条目"""
        try:
            commands_to_try = [
                ['grub-reboot', menu_path],
                ['/usr/sbin/grub-reboot', menu_path]
            ]
            
            success = False
            if os.geteuid() != 0:
                for cmd in commands_to_try:
                    try:
                        subprocess.run(['pkexec'] + cmd, check=True)
                        success = True
                        break
                    except (subprocess.CalledProcessError, FileNotFoundError):
                        try:
                            subprocess.run(['sudo'] + cmd, check=True)
                            success = True
                            break
                        except (subprocess.CalledProcessError, FileNotFoundError):
                            continue
            else:
                for cmd in commands_to_try:
                    try:
                        subprocess.run(cmd, check=True)
                        success = True
                        break
                    except (subprocess.CalledProcessError, FileNotFoundError):
                        continue
            
            if not success:
                raise RuntimeError("无法设置下次启动条目")
                
        except Exception as e:
            raise RuntimeError(f"设置下次启动条目失败: {str(e)}")

    def _clean_grub_entry_id(self, entry_id):
        """
        清理GRUB条目ID，确保它是一个有效的标识符
        """
        # 如果entry_id是复杂字符串（包含引号、特殊字符等），尝试提取简单版本
        if isinstance(entry_id, str):
            # 如果包含逗号，可能是完整菜单项，提取第一部分
            if ',' in entry_id:
                entry_id = entry_id.split(',')[0].strip()
            
            # 移除可能的引号
            entry_id = entry_id.strip('"\'')
            
            # 如果是数字，直接返回数字（GRUB索引方式）
            if entry_id.isdigit():
                return entry_id
            
            # 如果太长或者包含特殊字符，考虑使用索引方式
            if len(entry_id) > 100 or any(c in entry_id for c in ['--class', '$menuentry_id_option']):
                # 返回简单的索引方式
                return "0"  # 默认启动第一个条目
        
        return entry_id

    def update_grub_config(self):
        """更新GRUB配置"""
        try:
            # 首先尝试使用 update-grub (Debian/Ubuntu 系统)
            commands_to_try = [
                ['update-grub'],
                ['grub-mkconfig', '-o', self.grub_cfg_path],
                ['/usr/sbin/grub-mkconfig', '-o', self.grub_cfg_path]
            ]
            
            if os.geteuid() != 0:
                # 需要提权的情况
                for cmd in commands_to_try:
                    try:
                        subprocess.run(['pkexec'] + cmd, check=True)
                        return True
                    except (subprocess.CalledProcessError, FileNotFoundError):
                        try:
                            subprocess.run(['sudo'] + cmd, check=True)
                            return True
                        except (subprocess.CalledProcessError, FileNotFoundError):
                            continue  # 尝试下一个命令
                raise PermissionError("无法更新GRUB配置")
            else:
                # 已经是root权限
                for cmd in commands_to_try:
                    try:
                        subprocess.run(cmd, check=True)
                        return True
                    except (subprocess.CalledProcessError, FileNotFoundError):
                        continue  # 尝试下一个命令
                raise RuntimeError("无法更新GRUB配置")
                
        except Exception as e:
            print(f"更新GRUB配置失败: {e}")
            raise RuntimeError(f"无法更新GRUB配置: {e}")
    

class KernelManagementCLI:
    def __init__(self):
        self.boot_manager = BootManager(self)  # 传递自身实例给 BootManager
        
    def get_all_available_kernels(self):
        """获取所有可用的内核包"""
        try:
            # 使用包管理器检测
            package_manager = PackageManager()
            
            if package_manager.package_manager == 'apt':
                result = subprocess.run(['apt-cache', 'search', '^linux-image'], 
                                    capture_output=True, text=True, check=True)
                kernels = []
                for line in result.stdout.split('\n'):
                    if line.strip():
                        kernel_name = line.split()[0]
                        # 过滤掉包含'extra'和'dbg'的内核
                        if (kernel_name.startswith('linux-image-') and 
                            'extra' not in kernel_name and 
                            'dbg' not in kernel_name):
                            kernels.append(kernel_name)
                return kernels
            else:
                print(f"错误: 不支持的包管理器: {package_manager.package_manager}")
                return []
        except subprocess.CalledProcessError:
            print("错误: 无法获取可用内核列表")
            return []
        
    def get_installed_kernels(self):
        """列出所有已安装的内核"""
        try:
            # 使用包管理器检测
            package_manager = PackageManager()
            
            if package_manager.package_manager == 'apt':
                result = subprocess.run(['dpkg', '--list', 'linux-image-*'], 
                                      capture_output=True, text=True, check=True)
                kernels = []
                for line in result.stdout.split('\n'):
                    if line.startswith('ii'):
                        kernel_name = line.split()[1]
                        if kernel_name.startswith('linux-image-') and 'extra' not in kernel_name:
                            kernels.append(kernel_name)
                return kernels
            else:
                print(f"错误: 不支持的包管理器: {package_manager.package_manager}")
                return []
        except subprocess.CalledProcessError:
            print("错误: 无法获取已安装内核列表")
            return []
    
    def get_current_kernel(self):
        """获取当前正在使用的内核"""
        try:
            result = subprocess.run(['uname', '-r'], 
                                  capture_output=True, text=True, check=True)
            return result.stdout.strip()
        except subprocess.CalledProcessError:
            return None

    def list_installed_kernels(self):
        """列出已安装的内核"""
        print("正在获取已安装内核...")
        
        installed_kernels = self.get_installed_kernels()
        current_kernel = self.get_current_kernel()
        
        # 加载安装记录
        install_records = InstallRecordManager.load_records()
        
        # 加载默认启动和下次启动内核记录
        default_boot_kernel = DefaultBootManager.load_default_boot()
        next_boot_kernel = DefaultBootManager.load_next_boot()
        
        if not installed_kernels:
            print("未找到已安装的内核")
            return
        
        print(f"\n已安装内核列表 (总计: {len(installed_kernels)}):")
        print("-" * 80)
        
        for i, kernel in enumerate(installed_kernels, 1):
            # 根据不同包管理器提取版本号
            if kernel.startswith('linux-image-'):
                version = kernel.replace('linux-image-', '')
            else:
                version = kernel.replace('linux-', '')
            
            # 检查各种标记
            tags = []
            
            if current_kernel and (
                current_kernel == version or 
                version in current_kernel or 
                current_kernel in version):
                tags.append("当前内核")
            
            if next_boot_kernel and (
                next_boot_kernel == version or 
                version in next_boot_kernel or 
                next_boot_kernel in version):
                tags.append("下次启动")
            
            if default_boot_kernel and (
                default_boot_kernel == version or 
                version in default_boot_kernel or 
                default_boot_kernel in version):
                tags.append("默认启动")
            
            if version in install_records:
                tags.append("工具安装")
            
            tag_str = " [" + "][".join(tags) + "]" if tags else ""
            print(f"  {i:2d}. ✓ {kernel}{tag_str}")

    def list_available_kernels_paginated(self):
        """分页列出可安装的内核，支持搜索功能"""
        print("正在获取可用内核...")
        
        available_kernels = self.get_all_available_kernels()
        installed_kernels = self.get_installed_kernels()
        
        if not available_kernels:
            print("未找到可用内核")
            return
        
        # 提取已安装内核的版本号
        installed_versions = set()
        for kernel in installed_kernels:
            if kernel.startswith('linux-image-'):
                version = kernel.replace('linux-image-', '')
            else:
                version = kernel.replace('linux-', '')
            installed_versions.add(version)
        
        # 过滤出未安装的内核
        uninstall_kernels = []
        for kernel in available_kernels:
            if kernel.startswith('linux-image-'):
                version = kernel.replace('linux-image-', '')
            else:
                version = kernel.replace('linux-', '')
                
            if version not in installed_versions:
                uninstall_kernels.append(kernel)
        
        if not uninstall_kernels:
            print("没有可安装的内核")
            return
        
        # 搜索功能
        search_term = input("请输入搜索关键词（直接回车显示所有内核）: ").strip().lower()
        if search_term:
            filtered_kernels = [k for k in uninstall_kernels if search_term in k.lower()]
            if not filtered_kernels:
                print(f"未找到包含 '{search_term}' 的内核")
                return
            uninstall_kernels = filtered_kernels
            print(f"搜索到 {len(uninstall_kernels)} 个匹配的内核")
        
        # 分页显示
        page_size = 10
        total_pages = (len(uninstall_kernels) + page_size - 1) // page_size
        current_page = 1
        
        while True:
            start_idx = (current_page - 1) * page_size
            end_idx = min(start_idx + page_size, len(uninstall_kernels))
            
            print(f"\n可安装内核列表 (第 {current_page}/{total_pages} 页, 总计: {len(uninstall_kernels)}):")
            print("-" * 80)
            
            for i in range(start_idx, end_idx):
                kernel = uninstall_kernels[i]
                print(f"  {i+1:2d}. ○ {kernel}")
            
            # 显示分页导航
            nav_options = []
            if current_page > 1:
                nav_options.append("p: 上一页")
            if current_page < total_pages:
                nav_options.append("n: 下一页")
            nav_options.append("q: 返回")
            
            print(f"\n导航: {', '.join(nav_options)}")
            
            # 获取用户输入
            user_input = input("请选择操作: ").strip().lower()
            
            if user_input == 'p' and current_page > 1:
                current_page -= 1
            elif user_input == 'n' and current_page < total_pages:
                current_page += 1
            elif user_input == 'q':
                break
            else:
                print("无效选项，请重新选择")

    def install_kernel(self, kernel_name=None):
        """安装内核"""
        if not kernel_name:
            # 显示可用内核列表供用户选择
            available_kernels = self.get_all_available_kernels()
            installed_kernels = self.get_installed_kernels()
            
            print("可用内核列表:")
            uninstall_kernels = []
            for kernel in available_kernels:
                if kernel.startswith('linux-image-'):
                    version = kernel.replace('linux-image-', '')
                else:
                    version = kernel.replace('linux-', '')
                    
                installed_versions = [k.replace('linux-image-', '').replace('linux-', '') for k in installed_kernels]
                if version not in installed_versions:
                    uninstall_kernels.append(kernel)
            
            if not uninstall_kernels:
                print("没有可安装的内核")
                return
                
            for i, kernel in enumerate(uninstall_kernels):
                print(f"  {i+1}. {kernel}")
                
            try:
                choice = int(input("请选择要安装的内核编号: ")) - 1
                if choice < 0 or choice >= len(uninstall_kernels):
                    print("无效的选择")
                    return
                kernel_name = uninstall_kernels[choice]
            except ValueError:
                print("无效的输入")
                return
        
        # 确认安装
        confirm = input(f"确定要安装内核 {kernel_name} 吗？(y/N): ")
        if confirm.lower() != 'y':
            return
            
        print(f"开始安装内核: {kernel_name}")
        logger = CLILogger()
        logger.add_log(f"开始安装内核: {kernel_name}")
        
        try:
            # 使用包管理器安装
            package_manager = PackageManager()
            logger.add_log(f"检测到的包管理器: {package_manager.package_manager}")
            
            # 获取相关的包名
            base_version = kernel_name.replace('linux-image-', '').replace('linux-', '')
            package_names = package_manager.get_package_names(base_version, 'install')
            logger.add_log(f"将要安装的包: {', '.join(package_names)}")
            
            # 安装包
            installed_packages = package_manager.install_packages(package_names, logger)
                
            # 记录安装信息
            InstallRecordManager.add_record(base_version, installed_packages)
            logger.add_log(f"内核 {kernel_name} 安装完成，已记录安装信息")
            print(f"成功: 内核 {kernel_name} 安装完成")
            
        except Exception as e:
            logger.add_log(f"安装内核失败: {str(e)}")
            print(f"错误: 安装内核失败: {str(e)}")

    def remove_kernel(self, kernel_name=None):
        """移除内核"""
        installed_kernels = self.get_installed_kernels()
        current_kernel = self.get_current_kernel()
        
        if not kernel_name:
            # 显示已安装内核列表供用户选择
            if not installed_kernels:
                print("没有已安装的内核")
                return
                
            print("已安装内核列表:")
            removable_kernels = []
            for kernel in installed_kernels:
                # 检查是否有保护标记
                is_protected = False
                if kernel.startswith('linux-image-'):
                    version = kernel.replace('linux-image-', '')
                else:
                    version = kernel.replace('linux-', '')
                
                # 检查是否为当前内核
                if current_kernel and (
                    current_kernel == version or 
                    version in current_kernel or 
                    current_kernel in version):
                    is_protected = True
                
                if not is_protected:
                    removable_kernels.append(kernel)
                    print(f"  {len(removable_kernels)}. {kernel}")
            
            if not removable_kernels:
                print("没有可移除的内核")
                return
                
            try:
                choice = int(input("请选择要移除的内核编号: ")) - 1
                if choice < 0 or choice >= len(removable_kernels):
                    print("无效的选择")
                    return
                kernel_name = removable_kernels[choice]
            except ValueError:
                print("无效的输入")
                return
        
        # 确认移除
        confirm = input(f"确定要移除内核 {kernel_name} 吗？(y/N): ")
        if confirm.lower() != 'y':
            return
            
        print(f"开始移除内核: {kernel_name}")
        logger = CLILogger()
        logger.add_log(f"开始移除内核: {kernel_name}")
        
        try:
            # 使用包管理器移除
            package_manager = PackageManager()
            logger.add_log(f"检测到的包管理器: {package_manager.package_manager}")
            
            # 获取相关的包名
            base_version = kernel_name.replace('linux-image-', '').replace('linux-', '')
            package_names = package_manager.get_package_names(base_version, 'remove')
            logger.add_log(f"将要移除的包: {', '.join(package_names)}")
            
            # 移除包
            package_manager.remove_packages(package_names, logger)
            
            # 移除安装记录
            InstallRecordManager.remove_record(base_version)
            logger.add_log(f"内核 {kernel_name} 移除完成，已清除安装记录")
            print(f"成功: 内核 {kernel_name} 移除完成")
            
        except Exception as e:
            logger.add_log(f"移除内核失败: {str(e)}")
            print(f"错误: 移除内核失败: {str(e)}")

    def set_default_kernel(self, permanent=True):
        """设置默认启动内核"""
        self.boot_manager.set_default_kernel_with_recommended_method(permanent=permanent)

    def interactive_menu(self):
        """交互式菜单"""
        while True:
            print("\n=== Linux 内核管理工具 ===")
            print("1. 列出已安装内核")
            print("2. 列出可安装内核")
            print("3. 安装内核")
            print("4. 移除内核")
            print("5. 设置默认启动内核")
            print("6. 设置临时启动内核")
            print("7. 退出")
            print("=========================")
            
            try:
                choice = input("请选择操作 (1-7): ").strip()
                
                if choice == '1':
                    self.list_installed_kernels()
                elif choice == '2':
                    self.list_available_kernels_paginated()
                elif choice == '3':
                    self.install_kernel()
                elif choice == '4':
                    self.remove_kernel()
                elif choice == '5':
                    self.set_default_kernel(permanent=True)
                elif choice == '6':
                    self.set_default_kernel(permanent=False)
                elif choice == '7':
                    print("退出程序")
                    break
                else:
                    print("无效选项，请重新选择")
                    
            except KeyboardInterrupt:
                print("\n\n程序被用户中断")
                break
            except Exception as e:
                print(f"发生错误: {e}")

def main():
    cli = KernelManagementCLI()
    # 检查是否为root用户或具有sudo权限
    if os.geteuid() != 0:
        print("⚠️ 警告: 某些操作可能需要root权限")
    
    # 显示风险警告
    print("""
⚠️ 警告：内核管理操作具有高风险-使用本软件表示您已完全了解并同意所有风险。⚠️
- 在使用本工具前，请务必了解以下事项：
- 1. 安装或移除内核可能影响系统稳定性
- 2. 操作前请确保已备份重要数据
- 3. 建议在有经验的用户指导下操作
- 4. 软件作者对使用本软件产生的任何损失不承担任何责任
- 请务必阅读以上注意事项，并请务必保证您完全了解内核管理操作所带来的影响再继续使用该工具！
""")
    cli.interactive_menu()

if __name__ == "__main__":
    main()